import torch
import tqdm
from model import *
from dataset import *
import albumentations as A
from albumentations.pytorch import ToTensorV2
from Config import *
device = "cuda" if torch.cuda.is_available() else 'cpu'
print(device)

train_transformer = A.Compose([
    A.Resize(Config.img_size,Config.img_size),
    A.OneOf([
        A.Rotate(90,90,p=0.5)
    ]),
    A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.2),
    ToTensorV2(),
])


input_common = open(Config.train_input_path, 'rb').read()
gt_common = open(Config.train_gt_path, 'rb').read()
if __name__ == '__main__':
    dataset = Mydataset(input_common,gt_common,transformer=train_transformer,mode='train')
    dataloader = DataLoader(dataset,batch_size=Config.batch_size,shuffle=True,drop_last=True,pin_memory=True)
    # for i,(img,mask) in enumerate(dataloader):
    #     print(img,img.shape)
    model = Predictor()
    opt = torch.optim.Adam(model.parameters(), lr=Config.lr)
   # rnd = random.Random(100)
    criterion = nn.MSELoss(reduction='mean')
    for i in range(1,Config.train_epoch+1):
        losses = []
        model.train()
        for it,(img,mask) in enumerate(dataloader):
            img ,mask = img.to(device),mask.to(device)
            #print(img.dtype,img.shape,type(img))
            #image, mask = image.to(device), mask.to(device)
            opt.zero_grad()
            pred = model(img)
            loss = criterion(pred, mask)
            opt.zero_grad()
            loss.backward()
            opt.step()
            loss = float(loss.data.cpu().numpy())
            losses.append(loss)
            if it % 10 == 0:
                print('it', it, 'loss', loss, 'mean', np.mean(losses[-100:]))

    fout = open('model', 'wb')
    torch.save(model.state_dict(), fout)
    fout.close()
